import numpy as np
import time
from nes_py.wrappers import JoypadSpace
import gym_super_mario_bros
from gym_super_mario_bros.actions import SIMPLE_MOVEMENT, COMPLEX_MOVEMENT, RIGHT_ONLY
import math
from constants import *

class Unit:
	def __init__(self, eid, etype, state, world_x, world_y, vel_x, vel_y):
		self.eid = eid
		self.etype = etype
		self.state = state
		self.world_x = world_x
		self.world_y = world_y
		self.vel_x = vel_x
		self.vel_y = vel_y
		self.tile_x = -1
		self.tile_y = -1
		
	def get_tile_pos():
		pass
		
class Tile:
	def __init__(self, x, y, terr, unit):
		self.x = x # 0~15
		self.y = y # 0~14
		self.terr = terr
		self.unit = unit
		
def get_enemies(ram):
	enemies = []
	for slot in range(5):
		enemy = ram[0xF + slot]
		if enemy != 0:
			
			es = ram[addr_enemy_state + slot]
			#if es == 4 or es == 32 or es == 34 or es == 255:
			#	es = -1 # be killed
			#elif es == 196:
			#	es = 0 # fly koopa falling
			#elif es == 132:
			#	es = 1 # koopa tmp die
			
			et = ram[addr_enemy_type + slot]
			if et == 0:
				et = 4 # Green Koopa
			
			ex = ram[addr_enemy_page + slot] * 0x100 + ram[addr_enemy_x + slot]
			ey = ram[addr_enemy_y + slot]
			enemies.append({'type' : et, 'state' : es, 'x' : ex, 'y' : ey})
	return enemies

# get_x_position - Returns the current (horizontal) position
def get_x_position(ram):
	return ram[addr_curr_page] * 0x100 + ram[addr_curr_x]

# get_left_x_position - Returns number of pixels from left of screen
def get_left_x_position(ram):
	return (int(ram[addr_curr_x]) - int(ram[addr_left_x])) % 256

# get_y_position - Returns the current (vertical) position
def get_y_position(ram):
	return ram[addr_curr_y]

# get_y_viewport - Returns the current y viewport
# 1 = in visible viewport, 0 = above viewport, > 1 below viewport (i.e. dead)
def get_y_viewport(ram):
	return ram[addr_y_viewport]
	
def get_powerup_x(ram):
	return ram[addr_powerup_page] * 0x100 + ram[addr_powerup_x]

def get_powerup_y(ram):
	return ram[addr_powerup_y]

def get_tile_type(curr_x_position, left_x, box_x, box_y):
	x = curr_x_position - left_x + box_x + 112
	y = box_y + 96
	page = math.floor(x / 256) % 2
	
	sub_x = math.floor((x % 256) / 16)
	sub_y = math.floor((y - 32) / 16)
	curr_tile_addr = addr_tiles + page * 13 * 16 + sub_y * 16 + sub_x
	
	if (sub_y >= 13) or (sub_y < 0):
		return 0

	return env.ram[curr_tile_addr]

class GameData:
	def __init__(self, ram):
		self.ram = ram
		self.mario = None
		self.enemies = [None for _ in range(5)] # len=5
		self.tilemap = []
		for i in range(16):
			col = [Tile(i, j, 0, None) for j in range(15)]
			self.tilemap.append(col)
	
	def get_from_ram():
		# mario
		mario_x = get_x_position(self.ram)
		mario_y = get_y_position(self.ram)
		left_x = get_left_x_position(self.ram)


	